import os
import uuid

from model_layers import layer_dict
from util.console import console
import keras
from keras import Input, backend
from keras.layers import Dense
from keras.optimizers import Adam, RMSprop, SGD
from keras.utils import plot_model
from callbacks import keras_callbacks
from configuration.model_congig import ModelConfig
from log import log
from util.model_listutil import ModelListUtil

os.environ['TF_CPP_MIN_LOG_LEVEL'] = '2'


class Model:
    def __init__(self, conf: ModelConfig, inputs: Input):
        self.conf = conf

        self.opt = conf.OPTIMIZER(learning_rate=conf.LEARNING_RATE, decay=conf.DECAY)
        # self.opt = Adam(learning_rate=conf.LEARNING_RATE, decay=conf.DECAY)
        self.shape = (conf.IMAGE_HEIGHT, conf.IMAGE_WIDTH, conf.COLOR_CHANNELS)
        self.epoch = conf.EPOCHS
        self.inputs = inputs
        # self.opt = RMSprop(learning_rate=conf.LEARNING_RATE, decay=conf.DECAY)
        # self.opt = SGD(learning_rate=conf.LEARNING_RATE, decay=conf.DECAY, momentum=0.9, nesterov=True)
        print(self.opt)
    model = None
    history = None
    opt = None

    def create_by_list_real_cls_functional(self, model_list: list, models: list = None, last_dense: bool = True):
        """
        :param last_dense:
        :param models:
        :param model_list: [[keras.applications,*args][keras.applications,*args][keras.applications,*args]]
        :return:
        """
        inputs = self.inputs  # 保存原inputs
        ModelListUtil.check(model_list, models)
        for idx, model_list_line in enumerate(model_list):
            layer = layer_dict.get(model_list_line[0])
            self.inputs = layer().transfer(model_list_line, self.inputs, models)
        if last_dense:
            self.inputs = Dense(4, activation='softmax')(self.inputs)
        self.model = keras.Model(inputs, self.inputs)
        self.model.model_list = model_list
        self.model.conf = self.conf

    def compile(self):
        # self.model.aaa = 111
        self.model.compile(loss='categorical_crossentropy', optimizer=self.opt, metrics=["accuracy"])

    def train(self, x_train, y_train, x_test, y_test,
              batch_size: int = 64):
        log.warning(str(float(backend.get_value(self.model.optimizer.lr))))
        self.history = self.model.fit(x_train, y_train,
                                      batch_size=batch_size,
                                      epochs=self.epoch,
                                      validation_data=(x_test, y_test),
                                      shuffle=True,
                                      callbacks=keras_callbacks)

    @console
    def show(self, save: bool = True):
        self.model.summary(show_trainable=True, expand_nested=True)
        save_path = f'storage/model{uuid.uuid4().hex}.png'
        if save:
            plot_model(self.model,
                       show_shapes=True,
                       to_file=save_path,
                       show_layer_names=True,
                       show_layer_activations=True)
        self.model.save_path = save_path

    def save(self):
        # 保存模型和权重
        save_dir = os.path.join(os.getcwd(), 'saved_models')
        model_name = 'keras_garbage_trained_model.h5'

        if not os.path.isdir(save_dir):
            os.makedirs(save_dir)
        model_path = os.path.join(save_dir, model_name)
        self.model.save(model_path)
        log.info('Saved trained model at %s ' % model_path)

    def save_weights(self, path: str):
        self.model.save_weights(path)

    def evaluation(self, x_test, y_test, batch_size=64):
        # 评估训练模型
        scores = self.model.evaluate(x_test, y_test, verbose=1,batch_size=batch_size)
        log.info('Test loss:' + str(scores[0]))
        log.info('Test accuracy:' + str(scores[1]))

    def load(self, model_list: list, path: str):
        self.create_by_list_real_cls_functional(model_list)
        self.compile()
        self.model.load_weights(path)

    def predict(self, x):
        return self.model.predict(x)

    # def create_by_list_real_cls_functional(self, model_list: list, models: list = None, last_dense: bool = True,
    #                                         pop_num: int = 0):
    #     """
    #         :param model_list: [[keras.applications,*args][keras.applications,*args][keras.applications,*args]]
    #         :return:
    #         """
    #     ModelListUtil.check(model_list, models)
    #
    #     # inputs = Input(shape=self.shape)
    #     if model_list[0][0] == ModelLayers.Conv2D:
    #         filters = model_list[0][1]
    #         kernel_size = model_list[0][2]
    #         x = (Conv2D(filters,
    #     #                     (kernel_size, kernel_size),
    #                     padding='same',
    #                     # input_shape = shape
    #                     # input_shape=
    #                     ))(self.inputs)
    #     elif model_list[0][0] == ModelLayers.ConvBase:
    #         model_type = model_list[0][1]
    #         # 没写trainable默认False
    #         trainable = False if len(model_list[0]) < 3 else model_list[0][2]
    #         ConvBase = model_type(
    #             include_top=False,
    #             # input_shape=shape
    #         )
    #         ConvBase.trainable = trainable
    #         print(len(ConvBase.layers))
    #         trainable_num = 0 if len(model_list[0]) < 4 else model_list[0][3]
    #         if trainable_num < len(ConvBase.layers):
    #             for i in range(int(len(ConvBase.layers) - trainable_num)):
    #                 ConvBase.layers[i].trainable = False
    #         x = ConvBase(self.inputs)
    #     elif model_list[0][0] == ModelLayers.Concatenate:
    #         outputs = [model.outputs[0] for model in models]
    #         x = Concatenate()(outputs)
    #     elif model_list[0][0] == ModelLayers.Average:
    #         outputs = [model.outputs[0] for model in models]
    #         x = Average()(outputs)
    #     else:
    #         raise RuntimeError("check检查不完全")
    #     for model_item in model_list[1:]:
    #         if model_item[0] == ModelLayers.Conv2D:
    #             filters = model_item[1]
    #             kernel_size = model_item[2]
    #             # 这是业务逻辑check不检查
    #             if len(model_item) == 4 and model_item[3] == True:
    #                 x = Conv2D(filters, (kernel_size, kernel_size), padding='same', activation='relu')(x)
    #             else:
    #                 x = Conv2D(filters, (kernel_size, kernel_size), padding='same')(x)
    #
    #         elif model_item[0] == ModelLayers.Activation:
    #             activation = 'relu' if len(model_item) < 2 else model_item[1]
    #             x = Activation(activation)(x)
    #         elif model_item[0] == ModelLayers.MaxPooling2D:
    #             # 只写了MaxPooling2D没有写pool_size，用默认2
    #             pool_size = 2 if len(model_item) < 2 else model_item[1]
    #             x = MaxPooling2D(pool_size=(pool_size, pool_size))(x)
    #         elif model_item[0] == ModelLayers.Dropout:
    #             # 只写了Dropout没有写rate，用默认Config.DROP_OUT
    #             rate = Config.DROP_OUT if len(model_item) < 2 else model_item[1]
    #             x = Dropout(rate)(x)
    #         elif model_item[0] == ModelLayers.Flatten:
    #             x = Flatten()(x)
    #         elif model_item[0] == ModelLayers.BatchNormalization:
    #             x = BatchNormalizationV2()(x)
    #         elif model_item[0] == ModelLayers.Dense:
    #             units = model_item[1]
    #             # 没有写activation，用默认relu
    #             activation = 'relu' if len(model_item) < 3 else model_item[2]
    #             x = Dense(units, activation=activation)(x)
    #         else:
    #             raise RuntimeError("check检查不完全")
    #     if last_dense:
    #         x = Dense(4, activation='softmax')(x)
    #     # for _ in range(pop_num):
    #     #     x = x.node.parent_nodes.outputs
    #     self.model = keras.Model(self.inputs, x)
    #     self.model.model_list = model_list
    #     self.model.conf = self.conf
